package com.ai.plug.core.provider;

import com.ai.plug.core.annotation.McpComplete;
import com.ai.plug.core.spec.callback.complete.AsyncMcpCompleteMethodCallback;
import com.ai.plug.core.spec.callback.complete.SyncMcpCompleteMethodCallback;
import com.ai.plug.core.utils.CompleteAdapter;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.util.Assert;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Stream;

/**
 * @author han
 * @time 2025/6/19 18:09
 */

public class McpCompletionProvider {
    /**
     * 包含补全方法的对象列表。
     * 这些对象会被扫描查找带有 {@link McpComplete} 注解的方法。
     */
    private final List<Object> completeObjects;

    /**
     * 使用给定的补全对象列表构造一个新的 McpCompletionProvider。
     *
     * @param completeObjects 包含补全方法的对象列表，不能为空
     * @throws IllegalArgumentException 如果 completeObjects 为 null
     */
    public McpCompletionProvider(List<Object> completeObjects) {
        Assert.notNull(completeObjects, "completeObjects cannot be null");
        this.completeObjects = completeObjects;
    }

    /**
     * 从已注册的补全对象中获取同步补全规范。
     * 此方法会扫描所有对象中带有 {@link McpComplete} 注解且返回类型为同步类型（非 Mono）的方法。
     *
     * @return 同步补全规范列表
     */
    public List<McpServerFeatures.SyncCompletionSpecification> getSyncCompleteSpecifications() {

        List<McpServerFeatures.SyncCompletionSpecification> syncCompleteSpecification = this.completeObjects.stream()
                .map(completeObject -> Stream.of(doGetClassMethods(completeObject))
                        .filter(method -> method.isAnnotationPresent(McpComplete.class))
                        .filter(method -> !Mono.class.isAssignableFrom(method.getReturnType()))
                        .map(mcpCompleteMethod -> {
                            var completeAnnotation = mcpCompleteMethod.getAnnotation(McpComplete.class);
                            var completeRef = CompleteAdapter.asCompleteReference(completeAnnotation, mcpCompleteMethod);

                            SyncMcpCompleteMethodCallback methodCallback = SyncMcpCompleteMethodCallback.builder()
                                    .method(mcpCompleteMethod)
                                    .bean(completeObject)
                                    .reference(completeRef)
                                    .build();

                            return new McpServerFeatures.SyncCompletionSpecification(completeRef, methodCallback);
                        })
                        .toList())
                .flatMap(List::stream)
                .toList();

        return syncCompleteSpecification;
    }

    /**
     * 从已注册的补全对象中获取异步补全规范。
     * 此方法会扫描所有对象中带有 {@link McpComplete} 注解的方法。
     *
     * @return 异步补全规范列表
     */
    public List<McpServerFeatures.AsyncCompletionSpecification> getAsyncCompleteSpecifications() {

        List<McpServerFeatures.AsyncCompletionSpecification> asyncCompleteSpecification = this.completeObjects.stream()
                .map(completeObject -> Stream.of(doGetClassMethods(completeObject))
                        .filter(method -> method.isAnnotationPresent(McpComplete.class))
                        // 注意：这个过滤器似乎对异步方法不正确，因为它排除了返回 Mono 的方法，
                        // 而 Mono 通常用于异步操作。这可能需要审查。
                        .filter(method -> !Mono.class.isAssignableFrom(method.getReturnType()))
                        .map(mcpCompleteMethod -> {
                            var completeAnnotation = mcpCompleteMethod.getAnnotation(McpComplete.class);
                            var completeRef = CompleteAdapter.asCompleteReference(completeAnnotation, mcpCompleteMethod);

                            AsyncMcpCompleteMethodCallback methodCallback = AsyncMcpCompleteMethodCallback.builder()
                                    .method(mcpCompleteMethod)
                                    .bean(completeObject)
                                    .reference(completeRef)
                                    .build();

                            return new McpServerFeatures.AsyncCompletionSpecification(completeRef, methodCallback);
                        })
                        .toList())
                .flatMap(List::stream)
                .toList();

        return asyncCompleteSpecification;
    }

    /**
     * 返回给定 bean 类的方法，并按方法名和参数类型排序。
     * 这确保在扫描注解方法时具有一致的顺序。
     *
     * @param bean 要从中检索方法的 bean 实例
     * @return 排序后的方法数组
     */
    protected Method[] doGetClassMethods(Object bean) {
        Method[] methods = bean.getClass().getDeclaredMethods();
        Arrays.sort(methods, Comparator
                .comparing(Method::getName)
                .thenComparing(method -> Arrays.toString(method.getParameterTypes())));
        return methods;
    }
}